// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/serializers/presto_batch_vector_serializer.h>

#include <pollux/serializers/presto_serializer_estimation_utils.h>
#include <pollux/serializers/presto_serializer_serialization_utils.h>
#include <pollux/serializers/vector_stream.h>

namespace kumo::pollux::serializer::presto::detail {
    void PrestoBatchVectorSerializer::serialize(
        const RowVectorPtr &vector,
        const melon::Range<const IndexRange *> &ranges,
        Scratch &scratch,
        OutputStream *stream) {
        POLLUX_CHECK_NOT_NULL(vector, "Vector to serialize is null.");
        POLLUX_CHECK_NOT_NULL(stream, "Stream to serialize out to is null.");

#ifndef NDEBUG
        for (int i = 0; i < ranges.size(); i++) {
            POLLUX_CHECK_GE(ranges[i].begin, 0, "Invalid range at index {}", i);
            POLLUX_CHECK_LE(
                ranges[i].begin + ranges[i].size,
                vector->size(),
                "Invalid range at index {}",
                i);
        }
#endif

        SCOPE_EXIT {
            inUse.store(false);
        };

        POLLUX_CHECK(
            !inUse.exchange(true),
            "PrestoBatchVectorSerializer::serialize being called concurrently on the same object.");

        common::testutil::TestValue::adjust(
            "kumo::pollux::serializers::PrestoBatchVectorSerializer::serialize",
            this);

        const auto numRows = rangesTotalSize(ranges);
        const auto rowType = vector->type();
        const auto numChildren = vector->childrenSize();

        StreamArena arena(pool_);
        std::vector<VectorStream> streams;
        streams.reserve(numChildren);
        for (int i = 0; i < numChildren; i++) {
            streams.emplace_back(
                rowType->childAt(i),
                std::nullopt,
                vector->childAt(i),
                &arena,
                numRows,
                opts_);

            if (numRows > 0) {
                serializeColumn(vector->childAt(i), ranges, &streams[i], scratch);
            }
        }

        flushStreams(
            streams, numRows, arena, *codec_, opts_.minCompressionRatio, stream);
    }

    void PrestoBatchVectorSerializer::estimateSerializedSizeImpl(
        const VectorPtr &vector,
        const melon::Range<const IndexRange *> &ranges,
        vector_size_t **sizes,
        Scratch &scratch) {
        switch (vector->encoding()) {
            case VectorEncoding::Simple::FLAT:
                POLLUX_DYNAMIC_SCALAR_TYPE_DISPATCH_ALL(
                    estimateFlatSerializedSize,
                    vector->type_kind(),
                    vector.get(),
                    ranges,
                    sizes);
                break;
            case VectorEncoding::Simple::CONSTANT:
                POLLUX_DYNAMIC_TYPE_DISPATCH_ALL(
                    estimateConstantSerializedSize,
                    vector->type_kind(),
                    vector,
                    ranges,
                    sizes,
                    scratch);
                break;
            case VectorEncoding::Simple::DICTIONARY:
                POLLUX_DYNAMIC_TYPE_DISPATCH_ALL(
                    estimateDictionarySerializedSize,
                    vector->type_kind(),
                    vector,
                    ranges,
                    sizes,
                    scratch);
                break;
            case VectorEncoding::Simple::ROW: {
                if (!vector->may_have_nulls()) {
                    // Add the size of the offsets in the Row encoding.
                    for (int32_t i = 0; i < ranges.size(); ++i) {
                        *sizes[i] += ranges[i].size * sizeof(int32_t);
                    }

                    auto row_vector = vector->as<RowVector>();
                    auto &children = row_vector->children();
                    for (auto &child: children) {
                        if (child) {
                            estimateSerializedSizeImpl(child, ranges, sizes, scratch);
                        }
                    }

                    break;
                }

                std::vector<IndexRange> childRanges;
                std::vector<vector_size_t *> childSizes;
                for (int32_t i = 0; i < ranges.size(); ++i) {
                    // Add the size of the nulls bit mask.
                    *sizes[i] += bits::nbytes(ranges[i].size);

                    auto begin = ranges[i].begin;
                    auto end = begin + ranges[i].size;
                    for (auto offset = begin; offset < end; ++offset) {
                        // Add the size of the offset.
                        *sizes[i] += sizeof(int32_t);
                        if (!vector->is_null_at(offset)) {
                            childRanges.push_back(IndexRange{offset, 1});
                            childSizes.push_back(sizes[i]);
                        }
                    }
                }

                auto row_vector = vector->as<RowVector>();
                auto &children = row_vector->children();
                for (auto &child: children) {
                    if (child) {
                        estimateSerializedSizeImpl(
                            child,
                            melon::Range(childRanges.data(), childRanges.size()),
                            childSizes.data(),
                            scratch);
                    }
                }

                break;
            }
            case VectorEncoding::Simple::MAP: {
                auto map_vector = vector->as<MapVector>();
                std::vector<IndexRange> childRanges;
                std::vector<vector_size_t *> childSizes;
                expandRepeatedRanges(
                    map_vector,
                    map_vector->rawOffsets(),
                    map_vector->rawSizes(),
                    ranges,
                    sizes,
                    &childRanges,
                    &childSizes);
                estimateSerializedSizeImpl(
                    map_vector->mapKeys(), childRanges, childSizes.data(), scratch);
                estimateSerializedSizeImpl(
                    map_vector->mapValues(), childRanges, childSizes.data(), scratch);
                break;
            }
            case VectorEncoding::Simple::ARRAY: {
                auto array_vector = vector->as<ArrayVector>();
                std::vector<IndexRange> childRanges;
                std::vector<vector_size_t *> childSizes;
                expandRepeatedRanges(
                    array_vector,
                    array_vector->rawOffsets(),
                    array_vector->rawSizes(),
                    ranges,
                    sizes,
                    &childRanges,
                    &childSizes);
                estimateSerializedSizeImpl(
                    array_vector->elements(), childRanges, childSizes.data(), scratch);
                break;
            }
            case VectorEncoding::Simple::LAZY:
                estimateSerializedSizeImpl(
                    vector->as<LazyVector>()->loaded_vector_shared(),
                    ranges,
                    sizes,
                    scratch);
                break;
            default:
                POLLUX_UNSUPPORTED("Unsupported vector encoding {}", vector->encoding());
        }
    }
} // namespace kumo::pollux::serializer::presto::detail
